#include "optimal_binary_tree.h"

namespace NOptimalBinaryTreeHelper {

void FillOptimalBinaryTreeTables(size_t elemCount, const double* keyProbabilities, const double* missProbabilities, TTable<double>& e, TTable<double>& w, TTable<int>& root, std::ostream* dbgOutput) {
    for (size_t i = 0; i < elemCount + 1; ++i)
        w(i, i) = e(i, i) = missProbabilities[i]; // воображаемые узлы, обозначающие ненайденные элементы

    for (size_t l = 0; l < elemCount; ++l) {
        for (size_t i = 0; i < elemCount - l; ++i) {
            const size_t j = i + l + 1;
            assert(w(i, j - 1) >= 0.0); // != -1
            assert(j - 1 < elemCount);
            w(i, j) = w(i, j - 1) + keyProbabilities[j - 1] + missProbabilities[j];

            // Ищем оптимальную вершину для корня дерева [i; j)
            for (size_t r = i; r < j; ++r) {
                assert(e(i, r) >= 0.0); // != -1
                assert(e(r + 1, j) >= 0.0); // != -1
                const double tmp = e(i, r) + e(r + 1, j) + w(i, j);
                double& curr = e(i, j);
                if (curr < 0.0 || tmp < curr) {
                    curr = tmp;
                    root(i, j) = r;
                }
            }
        }
    }

    if (dbgOutput) {
        (*dbgOutput) << "e:\n" << e << std::endl;
        (*dbgOutput) << "w:\n" << w << std::endl;;
        (*dbgOutput) << "root:\n" << root << std::endl;
    }
}

} // namespace NOptimalBinaryTreeHelper
